USER = 'embed'
EMBED_TBL = 'kg_word_bank_m3e'

if '__main__' == __name__:
    from python_nlp.kg.neo4j.config import NEO4J_CONFIG_CBLUE_DETAIL
    from neo4j import GraphDatabase
    from PyCmpltrtok.common import sep, md5
    import pymongo as pm
    from PyCmpltrtok.util_mongo import mongo_upsert, VALUE, mongo_get
    from python_nlp.embed.cblue_text2mongo import sentence2embed, VER_INT, VERSION_COL, EMBED_COL, USERNAME
    import torch
    from transformers import AutoModel, AutoTokenizer
    from transformers import AutoConfig

    # 连接Mongodb
    mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
    mdb = mongo['CBLUE']

    # 加载文本BERT
    dev = torch.device(0)
    if 0:
        # model_name = 'bert-base-chinese'  # 用模型名，需联网，需翻墙
        model_name = 'moka-ai/m3e-base'  # 用模型名，需联网，需翻墙
    else:
        # model_name = r'C:\Users\peter\.cache\huggingface\hub\models--bert-base-chinese\snapshots\8d2a91f91cc38c96bb8b4556ba70c392f8d5ee55'  # 使用下载后的路径，不需联网
        model_name = r'D:\_const\wsl\my_github\m3e-base'
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(dev)

    # 确定最大token长度
    xlen01 = tokenizer.model_max_length
    config = AutoConfig.from_pretrained(model_name)
    xlen02 = config.max_position_embeddings
    max_len = min(xlen01, xlen02)

    def _main():

        print('Connecting neo4j ...')
        driver = GraphDatabase.driver(**NEO4J_CONFIG_CBLUE_DETAIL)
        print('Connected to neo4j.')
        with driver.session() as ss:
            for query in [
                "CREATE INDEX IF NOT EXISTS FOR (n:Node) ON (n.name)",
                "CREATE CONSTRAINT IF NOT EXISTS FOR (n:Node) REQUIRE (n.type, n.name) IS UNIQUE",
                "CREATE INDEX IF NOT EXISTS FOR (p:Predicate) ON (p.name)",
            ]:
                print(query)
                ss.run(query)
        print('Indexes and constraints OK.')

        # 确定名词
        with driver.session() as ss:
            print('Loading data ...')
            query = "MATCH (n:Node) RETURN DISTINCT n.name"
            print(query)
            xres = ss.run(query)
            print('Loaded.')

            for i, xel in enumerate(xres):
                xword = xel[0]
                # if i >= 20:
                #     break
                if i % 50 == 0:
                    print(i)
                print('.', end='')

                xmd5 = md5(xword)

                # 获取以有数据
                xdict = mongo_get(mdb, EMBED_TBL, USERNAME, xmd5, only_value=False)

                # 已经存在该句子的嵌入
                if xdict and xdict.get(EMBED_COL, None):
                    # print(f'Skip {xmd5} {xtext}')
                    continue

                # 做嵌入
                xembed = sentence2embed(dev, model, tokenizer, xword, max_seq_length=max_len)
                print('+', end='')

                # 入库
                mongo_upsert(mdb, EMBED_TBL, USER, xmd5, {
                    VALUE: xword,
                    EMBED_COL: xembed,
                    VERSION_COL: VER_INT,
                })
            print()
            print('Over')

    _main()